import numpy as np


def sample_prior(
    n_training_examples: int,
    low_boundary: np.ndarray | list,
    high_boundary: np.ndarray | list,
    seed=None
):
    "Sample parameter values from the prior distributions"

    RNG = np.random.default_rng(seed)

    n_thets = len(low_boundary)
    theta = RNG.uniform(
        low=low_boundary, high=high_boundary, size=(n_training_examples, n_thets)
    )

    return theta


def _dmcs_prior(
    n_training_examples: int,
    low_boundary: np.ndarray | list,
    high_boundary: np.ndarray | list,
    seed=None
):
    "Sample parameter values from the prior distributions"

    RNG = np.random.default_rng(seed)

    n_thets = len(low_boundary)
    theta = RNG.uniform(
        low=low_boundary, high=high_boundary, size=(n_training_examples, n_thets)
    )
    # generate congurency condition
    congurency = RNG.choice([1, 0], size=n_training_examples)
    # add congurency as last column in theta
    theta = np.column_stack((theta, congurency))

    return theta


def SSP_prior(
    n_training_examples=1,
    low_boundary=[0.2, 0.05, 0, 0.2, 0.01, 0.2],
    high_boundary=[5, 1, 5, 5, 5, 0.8],
    seed=None
):
    return _dmcs_prior(
        n_training_examples=n_training_examples,
        low_boundary=low_boundary,
        high_boundary=high_boundary,
        seed=seed
    )


def DMC_prior(
    n_training_examples=1,
    low_boundary=[0.2, 0.05, 0, 0.01, 1.1, 0.1, 0.2],
    high_boundary=[2.5, 1, 3, 2, 8, 2, 0.8],
    seed=None
):
    return _dmcs_prior(
        n_training_examples=n_training_examples,
        low_boundary=low_boundary,
        high_boundary=high_boundary,
        seed=seed
    )


def DDM_prior(
    n_examples: int,
    seed=None,
    trial_variability: bool = False,
):
    "Sample parameter values from the prior distributions"

    RNG = np.random.default_rng(seed)

    a,v,t,z = RNG.gamma(3, 1/3, size=n_examples), \
            RNG.normal(0, 2, size=n_examples), \
            RNG.gamma(12, 1/24, size=n_examples), \
            RNG.beta(10, 10, size=n_examples)
    if trial_variability:
        sv,st,sz = RNG.gamma(50, 1/100, size=n_examples), \
            RNG.gamma(20, 1/100, size=n_examples), \
            RNG.gamma(10, 1/100, size=n_examples), \

        return (a, v, t, z, sv, st, sz)
    else:
        return (a, v, t, z, 0, 0, 0)
